# -*- coding: utf-8 -*-
"""
Created on Fri Jun 10 10:31:19 2022

@author: qkrxp
"""
####################################################### 
# PROGRAM NAME : scEDI CALCULATOR                     #
# PROGRAMMER : Dr. Chang-Kyun Park (qkrxp2@gmail.com) # 
# REFERENCE : Park et al. 2022                        #
#######################################################
import numpy as np

###### TYPE YOUR CONFIGURATION ######
EDI_type = 'scEDI' # Choose the EDI type (scEDI or EDI)

input_file_path = 'EX_47108'
output_file_path = 'OUTPUT'

acc_num = 365    # Accumulation scale for dummy DS (default = 365 days)
Tday = 365       # Total days of 1 year (Leap years should be removed from input file before the calculation)
#####################################

####  MAIN PROGRAM (Do not modify) ####
read = np.loadtxt(input_file_path, dtype='str')
yr = np.array(read[:,0], dtype='int')
mo = np.array(read[:,1], dtype='int')
dy = np.array(read[:,2], dtype='int')
prec = np.array(read[:,3], dtype='float')

s_point = yr[0] + 29

MEAN_syr = np.zeros((len(prec)))
MEAN_eyr = np.zeros((len(prec)))
for i in range(len(prec)):
    if EDI_type == 'scEDI':
        if i >= acc_num:
            MEAN_syr[i] = yr[i] - 29
            MEAN_eyr[i] = yr[i]
            if yr[i] < s_point:
                MEAN_syr[i] = yr[0]
                MEAN_eyr[i] = yr[0] + 29  
    if EDI_type == 'EDI':
        if i >= acc_num:
            MEAN_syr[i] = yr[-1] - 29
            MEAN_eyr[i] = yr[-1]

EP = np.zeros((len(prec)))
MEP = np.zeros((len(prec)))
STD = np.zeros((len(prec))) 
CDD = np.zeros((len(prec)))
SD = acc_num
imo = 0
BA = 0; CN = 0    
for i in range(len(prec)):
    if i >= acc_num:
        C = MEAN_eyr[i] - MEAN_syr[i] + 1
        AvePrec = np.zeros((Tday))

        spt = int((MEAN_syr[i] - yr[0])*Tday)
        ept = int((MEAN_eyr[i] - yr[0])*Tday + Tday)
            
        dat = prec[spt:ept]
        for x in range(Tday):
            AvePrec[x] = np.mean(dat[x::Tday])        
        
        # For EDI: Equations (1), (3), (4) in Park et al. (2022)
        # For scEDI: Equations (13), (14) in Park et al. (2022)
        AP = 0; VAR = 0
        A = 0; B = 0
        for N in range(SD):
            AP = AP + prec[i-N]
            VAR = VAR + (AP/float(N+1))
            
            m = imo - N
            if m < 0:
                m = m + Tday
            A = A + AvePrec[m]
            B = B + (A / float(N+1))
        EP[i] = VAR
        MEP[i] = B
        
        A = EP[i] - MEP[i]
        if (A < 0) and (BA < 0):
            CN += 1
            CDD[i] = CN
        else:
            CN = 0
        BA = A
        
        imo += 1
        if imo == Tday:
            imo = 0 

imo = 0
for i in range(len(prec)):
    if i >= acc_num:    
        spt = int((MEAN_syr[i] - yr[0])*Tday)
        ept = int((MEAN_eyr[i] - yr[0])*Tday + Tday)
        
        dat = EP[spt:ept]
        STD[i] = np.std(dat[imo::Tday])
        
        imo += 1
        if imo == Tday:
            imo = 0    

print('EP, MEP, actual DS, and STD calculation complete')            
            
####  scEDI Calculation ####  
scEDI = np.zeros((len(prec)))
imo = 0
aline = ''
aline = aline + 'YEAR  MONTH  DAY    EDI     Precipitation' + '\n'
for i in range(len(prec)):
    if i >= acc_num: 
        
        if imo == Tday:
            imo = 0
        
        if CDD[i] == 0:
            CDDEP = EP[i]
            CDDMEP = MEP[i]
            CDDSTD = STD[i]
        else:
            SD = CDD[i] + acc_num
            C = MEAN_eyr[i] - MEAN_syr[i] + 1
            AvePrec = np.zeros((Tday))

            spt = int((MEAN_syr[i] - yr[0])*Tday)
            ept = int((MEAN_eyr[i] - yr[0])*Tday + Tday)
            
            dat = prec[spt:ept]
            for x in range(Tday):
                AvePrec[x] = np.mean(dat[x::Tday])
            
            # For EDI: Equations (6)-(8) in Park et al. (2022)
            # For scEDI: Equations (6), (16), (17) in Park et al. (2022)
            AP = 0; VAR = 0
            A = 0; B = 0
            for N in range(int(SD)):
                if (i-N) >= 0:
                    AP = AP + prec[i-N]
                    VAR = VAR + (AP/float(N+1))
                
                m = imo - N
                if m < 0:
                    for mx in range(999):
                        m = m + Tday
                        if m >= 0:
                            break
                A = A + AvePrec[m]
                B = B + (A / float(N+1))
            CDDMEP = B # MEP with CDD more than 0 is calculated
            CDDEP = VAR  # EP with CDD more than 0 is calculated
            
            # For EDI: Equations (9)-(11) in Park et al. (2022)
            # For scEDI: Equations (18)-(20) in Park et al. (2022)
            A = 0; B = 0
            for x in range(int(C)):
                YEARS = MEAN_syr[i] + x
                k = int((YEARS - yr[0])*Tday + imo)
                AP = 0; VAR = 0
                for N in range(int(SD)):
                    if (k-N) < 0:
                        break
                    AP = AP + prec[k-N]
                    VAR = VAR + (AP/float(N+1))
                B = VAR
                A = A + (B - CDDMEP)**2
            CDDSTD = np.sqrt(A / float(C))
        
        imo += 1
        
        scEDI[i] = (CDDEP - CDDMEP)/float(CDDSTD)
        if yr[i] >= (yr[0] + 30):
            aline = aline + str(yr[i])+ ' ' + str(mo[i]) + ' ' + str(dy[i]) + \
                ' ' + str(round(scEDI[i],2)) + ' ' + str(prec[i]) + '\n'

oidx = open(output_file_path,'w')
oidx.write(aline)
oidx.close()

print('scEDI calculation complete')

######### END
                
            
            
            
          
            
            
            
            
            
            
            
            
            
            
            
            
            
            
            
            
            
            
            
            
            
            
            
            
            
            